# This test script performs a number of image manipulation service on
# the given input image. It uses data decompisitioning to split up 
# the image into smaller parts that each surrogate may handle 
# individually.
#
# This scripts does a local split of the image, performs four image 
# operations using CF on the image fragments, and does a local merge.

from __future__ import with_statement
import sys
import scavenger
from scavenger import Scavenger
from time import sleep, time
from threading import Thread

# Check command line args.
if len(sys.argv) < 2:
    print 'Usage: imagetask.py imagefile [-i iterations] [-c cut-in-this-many-pieces]'
    sys.exit(1)

# Read in the image file.
with open(sys.argv[1], 'rb') as infile:
    image = infile.read()

# Check whether a number of test iterations is given.
iterations = 1
if '-i' in sys.argv:
    iterations = int(sys.argv[sys.argv.index('-i') + 1])

# Check whether a scheduler has been specified.
cut = 1
if '-c' in sys.argv:
    cut = int(sys.argv[sys.argv.index('-c') + 1])

# Sleep for a little while to make sure that we discover
# the available surrogates.
sleep(2.0)

def split_image(image, parts):
    from PIL import Image
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    output = {}
    width = pil_image.size[0]
    height = pil_image.size[1]
    lines_per_piece = height/parts
    for p in range(0, parts-1):
        image_fragment = pil_image.crop((0, p*lines_per_piece, width, (p+1)*lines_per_piece))
        sio = StringIO()
        image_fragment.save(sio, pil_image.format, quality=95)
        output['part%i'%p] = sio.getvalue()
    image_fragment = pil_image.crop((0, (parts-1)*lines_per_piece, width, height))
    sio = StringIO()
    image_fragment.save(sio, pil_image.format, quality=95)
    output['part%i'%(parts-1)] = sio.getvalue()
    return output

def merge_image(parts):
    from PIL import Image
    from StringIO import StringIO
    sorted_parts = []
    for partno in range(0, len(parts)):
        sio = StringIO(parts['part%i'%partno])
        image = Image.open(sio)
        sorted_parts.append(image)
    width = sorted_parts[0].size[0]
    height = 0
    for part in sorted_parts:
        height += part.size[1]
    image = Image.new('RGB', (width, height), None)
    current_height = 0
    for partno in range(0, len(sorted_parts)):
        image.paste(sorted_parts[partno], (0, current_height))
        current_height += sorted_parts[partno].size[1]
    sio = StringIO()
    image.save(sio, sorted_parts[0].format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', True)
def sharpen_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    factor = 1.0 + float(factor)
    sharpened_image = ImageEnhance.Sharpness(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', False)
def brightness_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Brightness(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', True)
def color_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Color(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', True)
def contrast_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Contrast(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

functions = [sharpen_batch, color_batch, contrast_batch, brightness_batch]
function_names = ['sharpen', 'color', 'contrast', 'brightness']

class Worker(Thread):
    def __init__(self, functions, image, x):
        super(Worker, self).__init__()
        self._functions = functions
        self._image = image
        self._x = x
        self.happy = True

    def run(self):
        global function_names
        try:
            for x in range(0, len(self._functions)):
                self._image = self._functions[x](self._image, 0.75, id="%s%i"%(function_names[x], self._x))
        except:
            self.happy = False    


# Run the test!
for x in range(0, iterations):
    start = time()
#    try:
    # Split the image into c parts.
    image_parts = split_image(image, cut)

    # Do the operations on the parts in parallel.
    workers = {}
    x = 1
    for key in image_parts.keys():
        w = Worker(functions, image_parts[key], x)
        w.start()
        workers[key] = w
        x += 1
    for key, worker in workers.items():
        worker.join()
        image_parts[key] = worker._image

    # Merge the parts back into a single image.
    image = merge_image(image_parts)

#    except:
#        print '-1',

    stop = time()

    # Check that all workers are happy.
    happy = True
    for worker in workers.values():
        if not worker.happy:
            happy = False
            break
    if not happy:
        print -1
    else:
        print stop - start

# Write the result for possible inspection.
with open('result.jpg', 'wb') as outfile:
    outfile.write(image)

scavenger.shutdown()
